package org.dja.nzw.discovery.zookeeper;

import org.apache.commons.io.IOUtils;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.x.discovery.*;
import org.apache.curator.x.discovery.details.JsonInstanceSerializer;
import org.dja.nzw.discovery.DiscoveryClient;
import org.dja.nzw.discovery.domain.InstanceDetails;
import org.dja.nzw.discovery.zookeeper.client.ClientFactory;

import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;

/**
 * Created by shuyun on 2016/7/18.
 */
public class ZKDiscoveryClient implements DiscoveryClient {
    private final String serverUrl;
    private static final String BASE_PATH = "/nzw/servers";

    private final AtomicReference<CuratorFramework> curatorFramework = new AtomicReference<CuratorFramework>();
    private static final JsonInstanceSerializer<InstanceDetails> serializer = new JsonInstanceSerializer<InstanceDetails>(InstanceDetails.class);

    private final Map<String, ServiceDiscovery<InstanceDetails>> serviceDiscoveries = new ConcurrentHashMap<>();
    private final Map<String, ServiceCache<InstanceDetails>> serviceCache = new ConcurrentHashMap<>();

    public ZKDiscoveryClient(String serverUrl) {
        this.serverUrl = serverUrl;
    }

    @Override
    public void startup() throws InterruptedException {
        curatorFramework.set(ClientFactory.create(serverUrl));
        curatorFramework.get().start();
    }

    @Override
    public void shutdown() throws InterruptedException {
        close();
        ClientFactory.destroy(curatorFramework.get());
    }

    @Override
    public List<InstanceDetails> applications(String appName) {
        return queryServiceInstance(getServiceCatch(appName));
    }

    @Override
    public List<InstanceDetails> applications(String appName, long version) {
        return queryServiceInstance(getServiceCatch(appName))
                .stream()
                .filter(ins -> version == ins.getVersion())
                .collect(Collectors.toList());
    }

    @Override
    public void registerService(InstanceDetails instance) {
        instance.setId(UUID.randomUUID().toString());
        instance.setVersion(System.currentTimeMillis());
        registerService(getServiceDiscovery(instance.getName()), instance);
    }

    @Override
    public void updateRegisterService(InstanceDetails instance) {
        instance.setVersion(System.currentTimeMillis());
        updateService(getServiceDiscovery(instance.getName()), instance);
    }

    @Override
    public void unregisterService(InstanceDetails instance) {

    }

    private String serverBasePath(){
        return BASE_PATH;
    }

    private ServiceDiscovery<InstanceDetails> getServiceDiscovery(String serviceName) {
        ServiceDiscovery<InstanceDetails> discovery = serviceDiscoveries.get(serviceName);
        if(null == discovery) {
            synchronized (serviceDiscoveries) {
                discovery = serviceDiscoveries.get(serviceName);
                if(null == discovery) {
                    try {
                        discovery = createServiceDiscovery();
                    } catch (Exception ex) {
                        throw new ExceptionInInitializerError(ex);
                    }
                    serviceDiscoveries.put(serviceName, discovery);
                }
            }
        }
        return discovery;
    }

    private ServiceDiscovery<InstanceDetails> createServiceDiscovery() throws Exception {
        ServiceDiscovery serviceDiscovery = ServiceDiscoveryBuilder.builder(InstanceDetails.class)
                .client(curatorFramework.get())
                .basePath(serverBasePath())
                .serializer(serializer)
                .build();

        serviceDiscovery.start();

        return serviceDiscovery;
    }

    private ServiceInstance<InstanceDetails> createServiceInstance(InstanceDetails instanceDetails) throws Exception{
        return ServiceInstance.<InstanceDetails>builder()
                .id(instanceDetails.getId())
                .name(instanceDetails.getName())
                .payload(instanceDetails)
                .uriSpec(new UriSpec("{scheme}://{address}:{port}"))
                .build();
    }

    private void registerService(ServiceDiscovery<InstanceDetails> serviceDiscovery, InstanceDetails instanceDetails) {
        try {
            ServiceInstance<InstanceDetails> instance = createServiceInstance(instanceDetails);
            serviceDiscovery.registerService(instance);
        } catch (Exception ex) {
            throw new ExceptionInInitializerError(ex);
        }
    }

    private void updateService(ServiceDiscovery<InstanceDetails> serviceDiscovery, InstanceDetails instanceDetails) {
        try {
            ServiceInstance<InstanceDetails> instance = createServiceInstance(instanceDetails);
            serviceDiscovery.registerService(instance);
        } catch (Exception ex) {
            throw new ExceptionInInitializerError(ex);
        }
    }

    private ServiceCache<InstanceDetails> getServiceCatch(String serviceName) {
        ServiceCache<InstanceDetails> cache = serviceCache.get(serviceName);
        if(null == cache) {
            synchronized (serviceCache) {
                cache = serviceCache.get(serviceName);
                if(null == cache) {
                    cache = initServiceCache(getServiceDiscovery(serviceName), serviceName);
                }
            }
        }

        return cache;
    }

    private ServiceCache initServiceCache(ServiceDiscovery discovery, String serviceName) {
        try {
            ServiceCache cache = createServiceCache(discovery, serviceName);
            serviceCache.put(serviceName, cache);
            return cache;
        } catch (Exception ex) {
            throw new ExceptionInInitializerError(ex);
        }
    }

    private ServiceCache createServiceCache(ServiceDiscovery discovery, String serviceName) throws Exception {
        ServiceCache cache = discovery.serviceCacheBuilder().name(serviceName).build();
        cache.start();

        return cache;
    }

    private List<InstanceDetails> queryServiceInstance(ServiceCache<InstanceDetails> cache) {
        return cache
                .getInstances()
                .stream()
                .map(ins -> ins.getPayload())
                .collect(Collectors.toList());
    }

    private void close(){
        serviceCache.forEach((name, cache) -> IOUtils.closeQuietly(cache));
        serviceCache.clear();
        serviceDiscoveries.forEach((name, discovery) -> IOUtils.closeQuietly(discovery));
        serviceDiscoveries.clear();
    }
}
